import numpy as np
import time
import pyopencl as cl
import verbose
import configuration as config
import pair 
import potential
import neighbor
import cell

ctx = cl.create_some_context()
que = cl.CommandQueue(ctx)
#opencl_index=[0,0]
#platform = cl.get_platforms()[opencl_index[0]]
#device = platform.get_devices()[opencl_index[1]]
#ctx = cl.Context([device]) 
#que = cl.CommandQueue(ctx)

v = verbose.Verbose()
v.set_verbose_level(0)

cfg = config.Configuration()
cfg.read_lammps_data('lj.data_32k.gz',datatype='pair')


#fact = 1.3 #enlarge box to reach packing fractor 0.2, compare to hoomd
fact = 1.0
cfg.box_l *= fact
cfg.box_hi *= fact
cfg.box_lo *= fact
#xyz = np.random.rand(cfg.n_atoms,4).astype(np.float32)
#xyz *= cfg.box_l
#cfg.r_t[:,0:3] = xyz[:,0:3]
#cfg.r_t[:,0:3] *= fact
#end random

cfg.init_device(ctx)
cfg.copy_to_device(que)
#cfg.n_atoms /= 10 #for test
print 'r_t[0:4]'
print cfg.r_t[0:4]


cel = cell.Cell(cfg,context=ctx,level=4,max_natoms_per_cell=16)
#cel = cell.Cell(cfg,context=ctx,level=5,max_natoms_per_cell=4)
t0 = time.time()
cel.cal_cell(que,(32000,),(64,),cfg)
cel.device_to_host(que)
cel.sort_on_host(cfg)
cfg.copy_to_device(que)
t1 = time.time()

print 'n_atoms_in_cell:'
print cel.n_atoms_in_cell
print 'atoms_in_cell:'
print cel.atoms_in_cell[0]
print cel.atoms_in_cell[1]
print 'cal-time:',t1-t0
print 'sorted order:'
print cfg.order
print 'sorted r_t[0:4]'
print cfg.r_t[0:4]


pots = potential.Potential(ctx)
neig = neighbor.Neighbor(cl_file='neighbor.cl',context=ctx)
print '--to-gen-grid'
t0 = time.time()
neig.create_grid_nei(que,cfg)
print 'create-nei(grid) time:', time.time()-t0
print 'min neighs/atom',np.min(neig.n_nei)
print 'max neighs/atom',np.max(neig.n_nei)
print 'ave neighs/atom',np.sum(neig.n_nei)/cfg.n_atoms
print 'neigh of atoms:'
print neig.nei[:,0]
print neig.nei[:,1]
print '-----'

lj = pair.Pair('pair_lj_cut','pair_lj_cut.cl',None,ctx)
pots.add_potential(lj)

nloop = 100
t0 = time.time()
for i in range(nloop):
    pots.cal_f_e_vir_all(que,cfg,neig)
cfg.get_from_device(que)
t1 = time.time()
#lj.read_parameter(None)
print 'f_e[0:10]:'
print cfg.f_e[0:10]
print nloop,' X cal_force time:',t1-t0
print 'aver force_en:',sum(cfg.f_e)/cfg.n_atoms
